import torch
from einops import rearrange
from torch import nn
from torch_geometric.nn import GCNConv
from torch_geometric.nn import MessagePassing
from torch_geometric.utils import add_self_loops, degree


class MyGCN(MessagePassing):
    def __init__(self, in_dim, out_dim):
        super().__init__(aggr="add")  # 加法聚合
        self.ff = nn.Linear(in_dim, out_dim)
        self.bias = nn.Parameter(torch.empty(out_dim))

    def forward(self, h, edge_index):
        # 增加自回路
        edge_index, _ = add_self_loops(edge_index, num_nodes=h.shape[0])

        # 计算h*seta
        h = self.ff(h)

        # 计算(D**(-1/2)) * A * (D**(-1/2))
        row, col = edge_index
        # 计算度张量  目标节点索引，节点数量，节点数据类型
        deg = degree(col, num_nodes=h.shape[0], dtype=h.dtype)
        deg_sqrt = deg ** -0.5
        # 度为0的倒数为inf，将它修改为0 ??这一步是否有必要，因为自环的度>=1
        deg_sqrt[deg_sqrt == float('inf')] = 0
        norm = deg_sqrt[row] * deg_sqrt[col]
        norm = rearrange(norm, "i->i ()")

        out = self.propagate(edge_index, h=h, norm=norm)  # 注意这里传入的不是h_i,是h
        out = out + self.bias
        return out

    def message(self, h_i, norm):
        return h_i * norm


class GCN(nn.Module):
    def __init__(self, in_dim, out_dim):
        super().__init__()
        self.conv1 = GCNConv(in_channels=in_dim, out_channels=16)
        self.conv2 = GCNConv(in_channels=16, out_channels=out_dim)

    def forward(self, x, edge_index):
        x = self.conv1(x, edge_index)
        out = self.conv2(x, edge_index)
        return out


if __name__ == "__main__":
    h = torch.randn(4, 8)  # 4个节点，每个节点维度为8的特征
    edge_index = torch.tensor([[0, 1],
                               [1, 0]])
    # 对边矩阵(2,E) E为边的数量，第一行为边的起点，第二行为边的目标节点

    my_gcn = MyGCN(in_dim=8, out_dim=16)
    h1 = my_gcn(h, edge_index)

    gcn = GCN(in_dim=8, out_dim=16)
    h2 = gcn(h, edge_index)

    pass
